import numpy as np
import matplotlib.pyplot as plt

def mesh_center_r(center_point, height, width, r=20):
    # 定义一个范围，即以 center_point 为中心，周围5个像素以内的范围
    range_x = np.arange(center_point[0] - r, center_point[0] + r+1)
    range_y = np.arange(center_point[1] - r, center_point[1] + r+1)
    # print(range_x, range_y)
    # 使用 np.meshgrid 生成所有可能的坐标组合
    xx, yy = np.meshgrid(range_x, range_y)

    # 将坐标组合展平并过滤掉超出图像边界的点
    valid_points = np.array([xx.flatten(), yy.flatten()]).T
    valid_points = valid_points[valid_points[:, 0] >= 0]
    valid_points = valid_points[valid_points[:, 0] < height]  # 假设 H 是图像高度
    valid_points = valid_points[valid_points[:, 1] >= 0]
    valid_points = valid_points[valid_points[:, 1] < width]  # 假设 W 是图像宽度

    # valid_points 现在包含了围绕 center_point 的周围5个像素以内的所有有效点
    # print(valid_points)  
    # 计算每个点到中心点的距离，并过滤掉超出半径的点
    distances = np.sqrt((valid_points[:, 0] - center_point[0]) ** 2 + (valid_points[:, 1] - center_point[1]) ** 2)
    valid_points = valid_points[distances <= r]
    
    return  valid_points

def mask_point_on_image(image, point, radius=20):
    
    H, W, C = image.shape
    
    valid_points = mesh_center_r(point, height=H, width=W, r=radius)
    # print(valid_points)

    # 将中心点周围的5个像素以内的有效点设置为白色
    center_mask = np.zeros_like(image, dtype=np.uint8)
    center_mask[valid_points[:, 0], valid_points[:, 1], :] = 1
    
    image *= (1 - center_mask)  # 白色区域设置为全黑
    # image[valid_points[:, 0], valid_points[:, 1], 0] = 255  # 白色区域设置为红色
    center_mask[:, :, :] = (255, 0, 0)  # 白色区域设置为红色
    image += center_mask  # 白色区域设置为红色
    
    return image

def test_center_r():
    H = 720  # 假设图像高度为 100
    W = 1280  # 假设图像宽度为 100

    # 生成一个全黑的彩色图像，三通道
    image = np.random.randint(0, 255, size=(H, W, 3), dtype=np.uint8)  # 假设图像为黑色

    center_point = (252, 300)  # 假设中心点为 (300, 400)
    
    image = mask_point_on_image(image, center_point, radius=15)
    # valid_points = mesh_center_r(center_point, height=H, width=W, r=20)
    # # print(valid_points)

    # # 将中心点周围的5个像素以内的有效点设置为白色
    # center_mask = np.zeros_like(image, dtype=np.uint8)
    # center_mask[valid_points[:, 0], valid_points[:, 1], :] = 1
    
    # image *= (1 - center_mask)  # 白色区域设置为全黑
    # # image[valid_points[:, 0], valid_points[:, 1], 0] = 255  # 白色区域设置为红色
    # center_mask[:, :, :] = (255, 0, 0)  # 白色区域设置为红色
    # image += center_mask  # 白色区域设置为红色

    # 画出中心点
    # image[center_point[0], center_point[1]] = [255, 255, 255]  # 白色
    # 设置打印选项以显示所有元素
    # np.set_printoptions(linewidth=np.inf, threshold=np.inf)
    # print(image)

    # 显示图像
    
    plt.imshow(image)
    plt.show()


def look_image(filename="elbow/video/00001.jpg"):
    import cv2
    
    img = cv2.imread(filename)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    H, W, C = img.shape
    
    point = (550, 1100)
    
    image = mask_point_on_image(img, point, radius=15)
    
    plt.imshow(image)
    plt.show()
    

if __name__ == '__main__':
    # test_center_r()
    look_image()
